#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
"""
@File    :   DirectConnection.py
@Time    :   2024-04-12 11:22
@Author  :   坐公交也用券
@Version :   1.0
@Contact :   faith01238@hotmail.com
@Homepage : https://liumou.site
@Desc    :   用于开启PVE网卡直通
"""
import re
from os import system
from sys import exit
from subprocess import getstatusoutput


class DirectConnection:
	def __init__(self):
		self.vt = False
		self.eth = ''
		self.eth_list = []

	def get_eth_list(self):
		"""
		获取网卡列表
		:return:
		"""
		cmd = 'ls /sys/class/net'
		res = getstatusoutput(cmd)
		if res[0] != 0:
			print("获取网卡列表失败")
			exit(2)

		# 使用正则表达式匹配物理网卡
		pattern = re.compile(r'\b(eth\w*|en\w*|eno\w*|em\w*)\b')
		physical_nics = pattern.findall(res[1])

		for nic in physical_nics:
			if nic.startswith(('eth', 'en', 'eno', 'em', 'wl')):
				self.eth_list.append(nic)

	def check_virtualization(self):
		with open('/proc/cpuinfo', 'r') as f:
			ci = f.read()

		if 'vmx' in ci or 'svm' in ci:
			self.vt = True
			return
		print("未开启VT-x/AMD-V虚拟化,请先进入BIOS开启")
		exit(2)

	def open_iommu(self):
		"""
		开启IOMMU
		:return:
		"""
		cmd = 'sed -i "s@GRUB_CMDLINE_LINUX_DEFAULT=.*@GRUB_CMDLINE_LINUX_DEFAULT=\"quiet intel_iommu=on\"@g" /etc/default/grub'
		system(cmd)
		if system("grep GRUB_CMDLINE_LINUX_DEFAULT /etc/default/grub") == 0:
			print("开启IOMMU成功")
		else:
			print("开启IOMMU失败")
			exit(3)

	def update_grub(self):
		"""

		:return:
		"""
		c = 'update-grub'
		if system(c) == 0:
			print("更新grub成功")
		else:
			print("更新grub失败")
			exit(4)

	def load_modules(self):
		"""
		加载模块
		:return:
		"""
		txt = "vfio\nvfio_iommu_type1\nvfio_pci\nvfio_virqfd"
		f = "/etc/modules"
		try:
			with open(f, 'a') as w:
				w.write(txt)
		except Exception as e:
			print("无法写入文件: ", f)
			print(e)
			exit(2)

	def update_initramfs(self):
		"""

		:return:
		"""
		c = 'update-initramfs -u -k all'
		if system(c) == 0:
			print("更新initramfs成功")
		else:
			print("更新initramfs失败")
			exit(5)

	def start(self):
		# self.get_eth_list()
		self.check_virtualization()
		self.open_iommu()
		self.update_grub()
		self.load_modules()
		self.update_initramfs()


if __name__ == "__main__":
	dc = DirectConnection()
	dc.start()
	# dc.get_eth_list()
